# -*- coding: utf-8 -*-
import os

# 设置环境变量以解决 joblib 警告
os.environ['JOBLIB_TEMP_FOLDER'] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'temp')
os.environ['JOBLIB_MULTIPROCESSING'] = '0'

import pandas as pd
import numpy as np
import joblib
from utils.log import Logger
import datetime
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score, precision_score, recall_score, \
    f1_score, roc_auc_score
import matplotlib.pyplot as plt
import seaborn as sns

plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15

# 定义预处理相关的文件名列表
PREPROCESSING_FILES = [
    'label_encoder.pkl',
    'onehot_encoder.pkl',
    'skewed_features.pkl',
    'scaler.pkl',
    'feature_names.pkl',
    'nominal_features.pkl'
]


def pred_feature_extract(data, logger):
    """
    预测数据特征提取，保持与模型训练时的特征列名一致
    1. 删除无用特征
    2. 添加交互特征
    3. 分箱处理
    4. 处理分类变量
    5. 处理偏态分布
    6. 标准化数值特征
    :param data: 输入数据
    :param logger: 日志对象
    :return: 处理后的特征数据
    """
    logger.info('=========开始特征提取==============')
    result = data.copy()

    # 1. 删除无用特征
    columns_to_drop = ['EmployeeNumber', 'Over18', 'StandardHours']
    result = result.drop(columns=columns_to_drop, axis=1)

    # 2. 添加交互特征
    logger.info('添加交互特征...')
    result['工作年限与年龄比'] = result['TotalWorkingYears'] / result['Age']
    result['晋升等待时间'] = result['YearsAtCompany'] - result['YearsSinceLastPromotion']
    result['薪资增长率'] = result['PercentSalaryHike'] / (result['YearsAtCompany'] + 1)
    result['工作满意度与绩效比'] = result['JobSatisfaction'] / result['PerformanceRating']
    result['工作投入度与满意度比'] = result['JobInvolvement'] / result['JobSatisfaction']
    result['工作生活平衡与满意度比'] = result['WorkLifeBalance'] / result['JobSatisfaction']

    # 3. 分箱处理
    logger.info('对连续变量进行分箱处理...')
    # 距离分组
    cut_labels = ['Near', 'Reasonable', 'Far']
    cut_bins = [-1, result['DistanceFromHome'].quantile(0.33),
                result['DistanceFromHome'].quantile(0.67),
                result['DistanceFromHome'].max() + 1]
    result['DistanceGroup'] = pd.cut(result['DistanceFromHome'], bins=cut_bins, labels=cut_labels)

    # 其他分箱
    result['年龄分组'] = pd.qcut(result['Age'], q=5, labels=['20-30', '30-40', '40-50', '50-60', '60+'])
    result['收入分组'] = pd.qcut(result['MonthlyIncome'], q=5, labels=['低', '中低', '中', '中高', '高'])
    result['工作年限分组'] = pd.qcut(result['TotalWorkingYears'], q=5,
                                     labels=['0-5年', '5-10年', '10-15年', '15-20年', '20年以上'])

    # 4. 处理分类变量
    logger.info('处理分类变量...')
    nominal_features = ['BusinessTravel', 'Department', 'EducationField',
                        'Gender', 'JobRole', 'MaritalStatus', 'OverTime',
                        'DistanceGroup', '年龄分组', '收入分组', '工作年限分组']

    # 加载OneHotEncoder
    encoder = joblib.load('../model/onehot_encoder.pkl')
    encoded = encoder.transform(result[nominal_features])
    encoded_df = pd.DataFrame(encoded, columns=encoder.get_feature_names_out(nominal_features), index=result.index)

    # 删除原始分类列
    result = result.drop(nominal_features, axis=1)

    # 合并编码后的特征
    result = pd.concat([result, encoded_df], axis=1)

    # 5. 处理偏态分布
    logger.info('处理偏态分布...')
    numeric_cols = result.select_dtypes(include=['float64', 'int64']).columns
    skewed_features = []
    for col in numeric_cols:
        if col != 'Attrition':  # 跳过目标变量
            skewness = result[col].skew()
            if abs(skewness) > 0.5:
                result[col] = np.log1p(result[col])
                skewed_features.append(col)
    logger.info(f"进行对数变换的特征: {skewed_features}")

    # 6. 标准化数值特征
    logger.info('标准化数值特征...')
    scaler = joblib.load('../model/scaler.pkl')
    numeric_cols = result.select_dtypes(include=['float64', 'int64']).columns
    if 'Attrition' in numeric_cols:
        numeric_cols = numeric_cols.drop('Attrition')
    result[numeric_cols] = scaler.transform(result[numeric_cols])

    # 7. 确保特征顺序与训练时一致
    feature_names = joblib.load('../model/feature_names.pkl')
    if 'Attrition' in feature_names:
        feature_names.remove('Attrition')

    # 确保所有特征都存在
    for feature in feature_names:
        if feature not in result.columns:
            result[feature] = 0

    # 按训练时的特征顺序排列
    result = result[feature_names]

    return result


def prediction_plot(data, model_name):
    """
    绘制预测结果可视化图
    :param data: 包含真实值和预测值的数据
    :param model_name: 模型名称
    """
    # 预测概率分布图
    plt.figure(figsize=(12, 6))
    plt.hist(data['预测概率'], bins=20)
    plt.title(f'{model_name} - 员工离职概率分布')
    plt.xlabel('离职概率')
    plt.ylabel('频数')
    plt.savefig(f'../data/fig/{model_name}_预测概率分布.png')
    plt.close()


def evaluate_predictions(y_true, y_pred, y_pred_proba, model_name, logger):
    """
    评估模型预测结果
    """
    logger.info(f'\n{model_name} 模型评估报告:')

    # 计算评估指标
    accuracy = accuracy_score(y_true, y_pred)
    precision = precision_score(y_true, y_pred)
    recall = recall_score(y_true, y_pred)
    f1 = f1_score(y_true, y_pred)
    auc = roc_auc_score(y_true, y_pred_proba)

    logger.info(f"准确率: {accuracy:.4f}")
    logger.info(f"精确率: {precision:.4f}")
    logger.info(f"召回率: {recall:.4f}")
    logger.info(f"F1分数: {f1:.4f}")
    logger.info(f"AUC: {auc:.4f}")

    # 打印分类报告
    logger.info("\n分类报告:")
    logger.info(classification_report(y_true, y_pred))

    # 绘制混淆矩阵
    plt.figure(figsize=(8, 6))
    cm = confusion_matrix(y_true, y_pred)
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
    plt.title(f'{model_name} 混淆矩阵')
    plt.ylabel('真实标签')
    plt.xlabel('预测标签')
    plt.savefig(f'../data/fig/{model_name}_test_confusion_matrix.png')
    plt.close()

    return {
        'Model': model_name,
        'Accuracy': accuracy,
        'Precision': precision,
        'Recall': recall,
        'F1 Score': f1,
        'AUC': auc
    }


def plot_model_comparison(results_df, logger):
    """绘制模型比较图表"""
    try:
        # 设置中文字体
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False

        # 创建图表
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))

        # 绘制Accuracy比较图
        sns.barplot(data=results_df, x='Model', y='Accuracy', ax=ax1)
        ax1.set_title('各模型准确率比较')
        ax1.set_ylim(0, 1)
        ax1.tick_params(axis='x', rotation=45)

        # 绘制F1 Score比较图
        sns.barplot(data=results_df, x='Model', y='F1 Score', ax=ax2)
        ax2.set_title('各模型F1分数比较')
        ax2.set_ylim(0, 1)
        ax2.tick_params(axis='x', rotation=45)

        plt.tight_layout()
        plt.savefig('model_comparison.png')
        plt.close()

        logger.info("模型比较图表已保存为 model_comparison.png")

    except Exception as e:
        logger.error(f"绘制模型比较图表时出错: {str(e)}")
        # 打印数据框的列名，以便调试
        logger.error(f"数据框的列名: {results_df.columns.tolist()}")
        raise


def plot_predictions(y_true, y_pred, model_name, logger):
    """绘制预测结果对比图"""
    try:
        # 设置中文字体
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False

        plt.figure(figsize=(10, 6))
        plt.scatter(y_true, y_pred, alpha=0.5)
        plt.plot([y_true.min(), y_true.max()], [y_true.min(), y_true.max()], 'r--', lw=2)
        plt.xlabel('实际值')
        plt.ylabel('预测值')
        plt.title(f'{model_name}预测结果对比')
        plt.tight_layout()
        plt.savefig(f'{model_name}_predictions.png')
        plt.close()

        logger.info(f"{model_name}预测结果对比图已保存")

    except Exception as e:
        logger.error(f"绘制预测结果对比图时出错: {str(e)}")
        raise


def plot_residuals(y_true, y_pred, model_name, logger):
    """绘制残差分布图"""
    try:
        # 设置中文字体
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False

        residuals = y_true - y_pred

        plt.figure(figsize=(10, 6))
        plt.scatter(y_pred, residuals, alpha=0.5)
        plt.axhline(y=0, color='r', linestyle='--')
        plt.xlabel('预测值')
        plt.ylabel('残差')
        plt.title(f'{model_name}残差分布')
        plt.tight_layout()
        plt.savefig(f'{model_name}_residuals.png')
        plt.close()

        logger.info(f"{model_name}残差分布图已保存")

    except Exception as e:
        logger.error(f"绘制残差分布图时出错: {str(e)}")
        raise


def find_best_model(results_df, logger):
    """
    根据AUC指标找出最优模型
    :param results_df: 包含所有模型评估结果的数据框
    :param logger: 日志对象
    :return: 最优模型名称
    """
    # 找出AUC最高的模型
    best_model = results_df.loc[results_df['AUC'].idxmax()]

    logger.info("\n最优模型分析 (基于AUC指标):")
    logger.info(f"最优模型: {best_model['Model']}")
    logger.info(f"AUC得分: {best_model['AUC']:.4f}")
    logger.info("\n其他评估指标:")
    for metric in ['Accuracy', 'Precision', 'Recall', 'F1 Score']:
        logger.info(f"{metric}: {best_model[metric]:.4f}")

    return best_model['Model']


def main():
    # 配置日志记录
    logfile_name = "predict_" + datetime.datetime.now().strftime('%Y%m%d%H%M%S')
    logger = Logger('../', logfile_name).get_logger()

    # 加载测试数据
    test_file = os.path.join('../data', 'test2.csv')
    test_data = pd.read_csv(test_file)
    logger.info(f'加载测试数据: {test_file}')

    # 特征提取
    X_test_processed = pred_feature_extract(test_data, logger)

    # 加载所有模型
    models = {}
    model_dir = '../model'
    for file in os.listdir(model_dir):
        if file.endswith('.pkl') and file not in PREPROCESSING_FILES:
            model_name = file.replace('.pkl', '')
            model_path = os.path.join(model_dir, file)
            try:
                model = joblib.load(model_path)
                # 验证是否是有效的模型对象
                if hasattr(model, 'predict_proba'):
                    models[model_name] = model
                else:
                    logger.warning(f"跳过文件 {file}，因为它不是一个有效的模型对象")
            except Exception as e:
                logger.error(f"加载模型 {model_name} 失败: {str(e)}")
                continue

    logger.info(f'加载了 {len(models)} 个模型')

    # 对每个模型进行预测和评估
    results = []

    for model_name, model in models.items():
        logger.info(f'\n使用 {model_name} 模型进行预测...')

        try:
            # 预测
            predictions = model.predict_proba(X_test_processed)[:, 1]
            y_pred = (predictions > 0.5).astype(int)

            # 保存预测结果
            pred_results = test_data.copy()
            pred_results['预测概率'] = predictions
            pred_results['预测结果'] = y_pred
            pred_results.to_csv(f'../data/{model_name}_predictions.csv', index=False)

            # 结果可视化
            prediction_plot(pred_results, model_name)

            # 如果有真实标签，进行评估
            if 'Attrition' in test_data.columns:
                y_true = test_data['Attrition']
                metrics = evaluate_predictions(y_true, y_pred, predictions, model_name, logger)
                results.append(metrics)
        except Exception as e:
            logger.error(f"模型 {model_name} 预测失败: {str(e)}")
            continue

    # 保存评估结果
    if results:
        results_df = pd.DataFrame(results)
        results_df.to_csv('../data/test_evaluation_results.csv', index=False)
        logger.info('\n所有模型的评估结果已保存到 test_evaluation_results.csv')

        # 绘制模型比较图
        plot_model_comparison(results_df, logger)

        # 找出最优模型
        best_model = find_best_model(results_df, logger)

        # 加载最优模型的预测结果
        best_model_pred = pd.read_csv(f'../data/{best_model}_predictions.csv')

        # 绘制最优模型的详细分析图
        plt.figure(figsize=(12, 6))
        sns.histplot(data=best_model_pred, x='预测概率', hue='Attrition', bins=30)
        plt.title(f'{best_model} - 预测概率分布（按真实标签）')
        plt.savefig(f'../data/fig/{best_model}_best_model_analysis.png')
        plt.close()

    print("预测完成，结果已保存到相应文件中")


if __name__ == '__main__':
    main()
